from matplotlib import pyplot as plt
from tensorflow.python.keras.models import Sequential
from tensorflow.python.keras.layers import Conv2D
from tensorflow.python.keras.layers import MaxPooling2D

from tensorflow.python.keras.layers import Dense
from tensorflow.python.keras.layers import GlobalAveragePooling2D
from tensorflow.python.keras.optimizer_v2.rmsprop import RMSprop

from keras.preprocessing.image import ImageDataGenerator

# 一个存放各个层的列表
Layers = [
	# 卷积层1
	Conv2D(filters=16, kernel_size=(3, 3), padding='same', activation='relu', input_shape=(128, 128, 3)),
	Conv2D(filters=16, kernel_size=(3, 3), padding='same', activation='relu'),
	# 池化层1
	MaxPooling2D(pool_size=(2, 2), strides=8),
	# 卷积层2
	Conv2D(filters=32, kernel_size=(3, 3), padding='same', activation='relu'),
	Conv2D(filters=32, kernel_size=(3, 3), padding='same', activation='relu'),
	# 池化层2
	MaxPooling2D(pool_size=(2, 2), strides=4),
	# 卷积层3
	Conv2D(filters=64, kernel_size=(3, 3), padding='same', activation='relu'),
	Conv2D(filters=64, kernel_size=(3, 3), padding='same', activation='relu'),
	# 池化层3
	MaxPooling2D(pool_size=(2, 2), strides=2),
	# 全局平均池化层
	GlobalAveragePooling2D(),
	# 全连接层
	Dense(64, activation='relu'),
	Dense(1, activation='sigmoid')

]


def LoadDataImage():
	train_data = ImageDataGenerator(
		rescale=1. / 255
	)

	test_data = ImageDataGenerator(rescale=1. / 255)

	train_set = train_data.flow_from_directory(
		'images/train_images',
		target_size=(128, 128),
		batch_size=5,
		class_mode='binary')

	test_set = test_data.flow_from_directory(
		'images/test_images',
		target_size=(128, 128),
		batch_size=20,
		class_mode='binary')
	return train_set, test_set


def test_accuracy(classifier, test_set, steps):
	num_correct = 0
	num_guesses = 0
	for i in range(steps):
		a = test_set.next()
		guesses = classifier.predict(a[0])
		correct = a[1]
		for index in range(len(guesses)):
			num_guesses += 1
			if round(guesses[index][0]) == correct[index]:
				num_correct += 1
	return num_correct, num_guesses


def Draw(history):
	acc = history.history['accuracy']
	val_acc = history.history['val_accuracy']
	loss = history.history['loss']
	val_loss = history.history['val_loss']

	epochs = range(len(acc))  # Get number of epochs

	plt.plot(epochs, acc, 'r', "Training Accuracy")
	plt.plot(epochs, val_acc, 'b', "Validation Accuracy")
	plt.title('Training and validation accuracy')
	plt.figure()
	plt.plot(epochs, loss, 'r', "Training Loss")
	plt.plot(epochs, val_loss, 'b', "Validation Loss")
	plt.title('Training and validation loss')
	plt.show()


def main():
	model = Sequential(Layers)
	model.compile(optimizer=RMSprop(lr=0.001), loss='binary_crossentropy', metrics=['accuracy'])
	train_set, test_set = LoadDataImage()

	history = model.fit_generator(
		train_set,
		epochs=50,
		validation_data=test_set

	)

	model.save('TrainedModel/model_new.h5')
	print("Model is saved!")
	num_correct, num_guesses = test_accuracy(model, test_set, 10)
	print("测了%d个图片,猜对了%d个" % (num_guesses, num_correct))
	Draw(history)


if __name__ == '__main__':
	main()
